import torch
import torch.nn as nn

class CNN(nn.Module):
    def __init__(self, input_dim=1280, hidden_dim=512, out_dim=128):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv1d(1, 2, kernel_size=5)
        self.fc1_input_dim = 2 * (input_dim - 4)
        self.fc1 = nn.Linear(self.fc1_input_dim, hidden_dim) 
        self.fc2 = nn.Linear(hidden_dim, out_dim)    

    def forward(self, x):
        # x.size [1000, input_dim]
        x = x.unsqueeze(1)  # x.size [1000, 1, input_dim]
        x = self.conv1(x)  # x.size [1000, 2, input_dim-4]
        x = torch.relu(x)
        x = x.view(x.size(0), -1)  # x.size [1000, 2*(input_dim-4)]
        x = self.fc1(x)  # x.size [1000, hidden_dim]
        x = torch.relu(x)
        x = self.fc2(x)  # x.size [1000, out_dim]
        x = torch.sigmoid(x)
        return x